import os, sys, logging
sys.path.append(".")
import matplotlib.pyplot as plt

import pyarma as pa
import numpy as np

from unifuncs import *
from expoint import *
from problems import *

def LYJ1D(u0, u1, tau, M, lb, ub, omega = 0.5, theta = 0):
    
    h = (ub - lb) / M
    r = tau / h
    
    A = pa.zeros(M, M)
    B = pa.zeros(M, 1)
    
    for i in range(1, M):
        
        if i < M - 1:
            
            A[i, i+1] = (1 - omega) * r / 2 + r * u1[i] / 6 + r * u1[i+1] / 6 - (1 + theta) / (h ** 2)
            B[i] += (-omega * r - (2 * theta) / (h ** 2)) * u1[i+1] + (-(1 - omega) * r / 2 - r * u1[i] / 6 - r * u1[i+1] / 6 + (theta - 1) / (h ** 2)) * u0[i+1]
            
        A[i, i] = 1 + (2 + 2 * theta) / (h ** 2)
        B[i] += 4 * theta * u1[i] / (h ** 2) + (1 - (2 * theta - 2) / h ** 2) * u0[i]
        
        if i > 1:
            
            A[i, i-1] = - (1 - omega) * r / 2 - r * u1[i] / 6 - r * u1[i-1] / 6 - (1 + theta) / (h ** 2)
            B[i] += (omega * r - 2 * theta / (h ** 2)) * u1[i-1] + ((1 - omega) / 2 + r * u1[i] / 6 + r * u1[i-1] / 6 + (theta - 1) / (h ** 2)) * u0[i - 1]
            
    u2 = pa.zeros(M, 1)
    u2[0] = u1[0] + tau
    u2[1:M-1, 0] = pa.solve(A[1:M-1, 1:M-1], B[1:M-1, 0])
    
    return u2

def EgLYJ(M, lb, ub, tau = 0.1, T = 20, c=0.1, delta=1, x0 = 0, Launcher = RK.StrehmelStrong2, omega = 0.5, theta = 0):
    
    h = (ub - lb) / M

    u_real = lambda t: rlw1d.SWave1(M, lb, ub, t, c, delta, x0)

    F, J, g, A = rlw1d.Finite_Difference(M, lb, ub, delta)

    t = [0]
    tn = np.int32(T / tau + 1)
    u = [u_real(0)]
    e = [pa.zeros(M, 1)]

    logger = logging.getLogger(__name__)
    logger.setLevel(level = logging.INFO)
    handler = logging.FileHandler("./logs/logs.txt")
    handler.setLevel(level = logging.INFO)
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    I = rlw1d.Invariant(c, delta)

    logger.info("Simulation started: Method = SHFDM-EIRK with scheme %s, xRange = [%.2lf, %.2lf], h = %.3lf, tau = %.3lf, amplitude = %.2lf, Invariants_Analytic = (%.9lf, %.9lf, %.9lf)" % ("LYJ_FD", lb, ub, h, tau, 3 * c, I[0], I[1], I[2]))
    print("Simulation started: Method = SHFDM-EIRK with scheme %s, xRange = [%.2lf, %.2lf], h = %.3lf, tau = %.3lf, amplitude = %.2lf, Invariants_Analytic = (%.9lf, %.9lf, %.9lf)" % ("LYJ_FD", lb, ub, h, tau, 3 * c, I[0], I[1], I[2]))

    LClist = np.zeros((tn, 6), dtype = np.float128)

    for k in range(2):
        if k > 0:
            t.append(t[k - 1] + tau)
            u.append(Launcher(u[k-1], tau, A, g))
        LClist[k, 0] = t[k]
        LClist[k, 1:3] = rlw1d.get_Errors(M, lb, ub, u[k], u_real)[1:3]
        LClist[k, 3:6] = rlw1d.get_Invariant(M, lb, ub, u[k], delta)
        
        logger.info('(Iter = %d, CPU Time = %s) Time = %.2lf, L2, Linf, M, P, E = %.8e\t&\t%.8e\t&\t%.9lf\t&\t%.9lf\t&\t%.9lf' % (k, "----------", LClist[k, 0], LClist[k, 1], LClist[k, 2], LClist[k, 3], LClist[k, 4], LClist[k, 5]))
        print(('(Iter = %d, CPU Time = %s) Time = %.2lf, L2, Linf, M, P, E = %.8e\t&\t%.8e\t&\t%.9lf\t&\t%.9lf\t&\t%.9lf' % (k, "----------", LClist[k, 0], LClist[k, 1], LClist[k, 2], LClist[k, 3], LClist[k, 4], LClist[k, 5])))

    timer = pa.wall_clock()

    for k in range(2, tn):

        t.append(t[k - 1] + tau)
        timer.tic()
        u.append(LYJ1D(u[k-2], u[k-1], tau, M, lb, ub, omega, theta))
        timecost = timer.toc()
        LClist[k, 0] = t[k]
        LClist[k, 1:3] = rlw1d.get_Errors(M, lb, ub, u[k], u_real)[1:3]
        LClist[k, 3:6] = rlw1d.get_Invariant(M, lb, ub, u[k], delta)
        logger.info('(Iter = %d, CPU Time = %s) Time = %.2lf, L2, Linf, M, P, E = %.8e\t&\t%.8e\t&\t%.9lf\t&\t%.9lf\t&\t%.9lf' % (k, str(timecost), LClist[k, 0], LClist[k, 1], LClist[k, 2], LClist[k, 3], LClist[k, 4], LClist[k, 5]))
        print(('(Iter = %d, CPU Time = %s) Time = %.2lf, L2, Linf, M, P, E = %.8e\t&\t%.8e\t&\t%.9lf\t&\t%.9lf\t&\t%.9lf' % (k, str(timecost), LClist[k, 0], LClist[k, 1], LClist[k, 2], LClist[k, 3], LClist[k, 4], LClist[k, 5])))
        
    fig, ax = plt.subplots()
    ax.set(xlabel = r"$x$", ylabel = r"$u$")
    x = np.linspace(lb + h, ub - h, M-1)
    yb, ym, ye = np.array(u[0][1:M-1, 0]), np.array(u[(tn - 1) // 2][1:M-1, 0]), np.array(u[tn - 1][1:M-1, 0])
    ax.plot(x, yb, label = r"$t=%.1lf$" % (0), color = "k")
    ax.plot(x, ym, label = r"$t=%.1lf$" % (t[(tn - 1) // 2]), color = "b")
    ax.plot(x, ye, label = r"$t=%.1lf$" % (t[tn - 1]), color = "r")
    ax.legend()
    plt.savefig("./outputs/result_method=%s_c=%.2lf_delta=%.2lf_xRange=[%.1lf,%.1lf]_h=%.3lf,tau=%.3lf.png" % ("LYJ_FD", c, delta, lb, ub, h, tau))
    
    np.savetxt("./outputs/result_method=%s_c=%.2lf_delta=%.2lf_xRange=[%.1lf,%.1lf]_h=%.3lf,tau=%.3lf.csv" % ("LYJ_FD", c, delta, lb, ub, h, tau), LClist, header = 'Time,L2_Error,Li_Error,M,P,E', delimiter = ',')
    logger.info("Simulation ended. MaxDiff: L2 = %.8e, Linf = %.8e\n" % (np.max(LClist[:, 1]), np.max(LClist[:, 2])))
    print("Simulation ended. MaxDiff: L2 = %.8e, Linf = %.8e" % (np.max(LClist[:, 1]), np.max(LClist[:, 2])))

    logger.removeHandler(handler)
    return u, LClist

if __name__ == '__main__':
    
    EgLYJ(1000, -40, 60, tau = 0.1, T = 20, c=0.1, delta=1, x0 = 0, Launcher = RK.StrehmelStrong2, omega = 1.1, theta = 0)